WORKFLOW_DIR=$(dirname "$(realpath "$0")")
cd "$WORKFLOW_DIR"/../../scripts
pwd

SCRIPT_DIR=$(dirname "$(realpath "$0")")
PROJECT_DIR=$(realpath "$SCRIPT_DIR"/../)
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

mkdir -p ${PROJECT_DIR}/outputs

echo "start sh run_llm_pretrain.sh----------------------------"
sh run.sh \
    --train_stage "PT" \
    --model_size 1B-dense \
    --learning_rate 4e-4 \
    --num_warmup_steps 200 \
    --num_train_epochs 10 \
    --max_seq_len 512 \
    --data_path "${PROJECT_DIR}/resource/dataset/minimind2/pretrain_1k_s512.jsonl" \
    --images_path None \
    --images_lmdb_path None \
    --load_ckpt_path "" \
    --save_dir ${PROJECT_DIR}/outputs/PT \
    --train_micro_batch_size_per_gpu 1 \
    --gradient_accumulation_steps 1 \
    --zero_stage 1 \

wait
echo "start sh run_llm_sft1.sh----------------------------"
sh run.sh \
    --train_stage "SFT"  \
    --model_size 1B-dense  \
    --learning_rate 5e-6 \
    --num_warmup_steps 50 \
    --num_train_epochs 1 \
    --max_seq_len 512 \
    --data_path ${PROJECT_DIR}/resource/dataset/minimind2/sft_1k_s512.jsonl \
    --images_path None \
    --images_lmdb_path None \
    --load_ckpt_path ${PROJECT_DIR}/outputs/PT/PT.pth \
    --save_dir ${PROJECT_DIR}/outputs/SFT \
    --train_micro_batch_size_per_gpu 1 \
    --gradient_accumulation_steps 1 \
    --zero_stage 1 \

wait

echo "start sh run_llm_sft2.sh----------------------------"
sh run.sh \
    --train_stage "SFT"  \
    --model_size 1B-dense  \
    --learning_rate 5e-6 \
    --num_warmup_steps 50 \
    --num_train_epochs 1 \
    --max_seq_len 2048 \
    --data_path ${PROJECT_DIR}/resource/dataset/minimind2/sft_1k_s2048.jsonl \
    --images_path None \
    --images_lmdb_path None \
    --load_ckpt_path ${PROJECT_DIR}/outputs/SFT/SFT.pth \
    --save_dir ${PROJECT_DIR}/outputs/SFT_2 \
    --train_micro_batch_size_per_gpu 1 \
    --gradient_accumulation_steps 1 \
    --zero_stage 1 \

wait
echo "start sh run_llm_dpo.sh----------------------------"
sh run.sh \
    --train_stage "DPO"  \
    --model_size 1B-dense  \
    --learning_rate 5e-8 \
    --num_warmup_steps 50 \
    --num_train_epochs 1 \
    --max_seq_len 2048 \
    --data_path ${PROJECT_DIR}/resource/dataset/minimind2/dpo_1k_s512.jsonl \
    --images_path None \
    --images_lmdb_path None \
    --load_ckpt_path ${PROJECT_DIR}/outputs/SFT_2/SFT.pth \
    --save_dir ${PROJECT_DIR}/outputs/DPO \
    --train_micro_batch_size_per_gpu 1 \
    --gradient_accumulation_steps 1 \
    --zero_stage 1 \

wait

echo "start sh run_llm_distillation.sh----------------------------"
sh run.sh \
    --train_stage "Distillation"  \
    --model_size 1B-dense  \
    --learning_rate 5e-6 \
    --num_warmup_steps 50 \
    --num_train_epochs 1 \
    --max_seq_len 1024 \
    --data_path ${PROJECT_DIR}/resource/dataset/minimind2/r1_1k_s1024.jsonl \
    --images_path None \
    --images_lmdb_path None \
    --load_ckpt_path ${PROJECT_DIR}/outputs/DPO/DPO.pth \
    --save_dir ${PROJECT_DIR}/outputs/Distillation \
    --train_micro_batch_size_per_gpu 1 \
    --gradient_accumulation_steps 1 \
    --zero_stage 1 \
